import tensorflow as tf
import pickle
import os
from sklearn.preprocessing import LabelEncoder

# 模型路径配置
MODEL_PATH = 'model/final_model.keras'
CLASS_MAPPING_PATH = 'model/class_mapping.pkl'


def load_model():
    """加载训练好的模型"""
    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(f"模型文件不存在: {MODEL_PATH}")

    try:
        # 导入TCN层定义，解决加载包含TCN层的模型时的错误
        from tcn import TCN
        model = tf.keras.models.load_model(MODEL_PATH, custom_objects={'TCN': TCN})
        print(f"模型加载成功: {MODEL_PATH}")
        return model
    except Exception as e:
        raise Exception(f"模型加载失败: {str(e)}")


def get_class_mapping():
    """获取类别映射关系"""
    # 实际应用中，应该在训练时保存class_mapping，这里简化处理
    if os.path.exists(CLASS_MAPPING_PATH):
        with open(CLASS_MAPPING_PATH, 'rb') as f:
            return pickle.load(f)

    # 如果没有保存的映射文件，这里应该根据实际情况返回正确的映射
    # 注意：这只是示例，实际应用中需要替换为真实的类别映射
    return {0: '晴', 1: '阴', 2: '雨', 3: '雪'}  # 示例映射


def save_class_mapping(class_mapping):
    """保存类别映射关系（训练时使用）"""
    with open(CLASS_MAPPING_PATH, 'wb') as f:
        pickle.dump(class_mapping, f)
